import os
import cv2
import sys
from .transformations import resize_transform_basic  #, rga_train_transform, rga_test_transform
import numpy as np
import pandas as pd
import random
import cv2
import PIL
from pathlib import Path
import torch
from torch.utils.data import DataLoader, Dataset, sampler
from utils.helper import read_xml, find_nodes, change_node_text, indent, write_xml

class T7_Mask_Dataset(Dataset):
    '''
    T7_Mask_Dataset
    '''
    def __init__(self, dataframe, transform=None):
        self.df = dataframe
        if transform is None:
            self.seg_transforms = resize_transform_basic()
        else:
            self.seg_transforms = transform

    def __len__(self):
        return self.df.shape[0]

    def __getitem__(self, idx):

        imgPath = self.df["image"].iloc[idx]
        label = self.df["label"].iloc[idx]
        try:
            image = cv2.imread(imgPath)
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        except IOError:
            print(imgPath)
            raise IOError("The image path is incorrect, please check the image path.")

        aug = self.seg_transforms(image=image)

        if os.path.isfile(label):
            label = Path(label).parts[-2]
            
        return aug['image'], label, Path(imgPath).stem, imgPath


class T10_Small_Mask_Dataset(Dataset):
    '''
    T7_Mask_Dataset
    '''

    def __init__(self, dataframe, transform=None, GRID_NUM=2):
        self.df = dataframe
        self.GRID_NUM = GRID_NUM
        if transform is None:
            self.seg_transforms = resize_transform_basic()
        else:
            self.seg_transforms = transform

    def __len__(self):
        return self.df.shape[0]

    def __getitem__(self, idx):

        imgPath = self.df["image"].iloc[idx]
        label = self.df["label"].iloc[idx]
        try:
            image = cv2.imread(imgPath)
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            H, W, C = image.shape
        except IOError:
            print(imgPath)
            raise IOError("The image path is incorrect, please check the image path.")

        quarter_h = int(H // self.GRID_NUM)
        quarter_w = int(W // self.GRID_NUM)

        aug_images = []
        for i in range(self.GRID_NUM):
            for j in range(self.GRID_NUM):
                image_roi = image[i * quarter_h: (i + 1) * quarter_h, j * quarter_w: (j + 1) * quarter_w, ...]
                aug = self.seg_transforms(image=image_roi)
                aug_images.append(aug['image'])

        # aug = self.seg_transforms(image=image)

        if os.path.isfile(label):
            label = Path(label).parts[-2]

        return aug_images, label, Path(imgPath).stem, imgPath

def box_fix_crop(img, bbox, size=224, pad_value=0):
    h, w, _ = img.shape
    if isinstance(size, int):
        size = (size, size)  # h, w

    # img add border
    # pad = int(max(size) * 3 // 4)
    # top, bottom, left, right = pad, pad, pad, pad
    # img_pad = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, None, pad_value)

    xmin, ymin, xmax, ymax = bbox
    # ymin = int(cy-int(box_h/2)+1)
    # xmin = int(cx-int(box_w/2)+1)
    # ymax = int(cy+int(box_h/2)+1)
    # xmax = int(cx+int(box_w/2)+1)
    center_x = (xmax + xmin) / 2.
    center_y = (ymax + ymin) / 2.
    defect_w, defect_h = xmax-xmin, ymax-ymin

    if defect_w/defect_h > size[1] / size[0]:
        new_defect_h = size[0] * defect_w / size[1]
        new_defect_w = defect_w
    else:
        new_defect_w = size[1] * defect_h / size[0]
        new_defect_h = defect_h
    
    if new_defect_h < size[0]:
        new_defect_h, new_defect_w = size[0], size[1]

    new_x_min = max(int(center_x-int(new_defect_w/2)+1), 0)
    new_y_min = max(int(center_y-int(defect_h/2)+1), 0)
    new_x_max = min(int(center_x+int(new_defect_w/2)+1), w)
    new_y_max = min(int(center_y+int(defect_h/2)+1), h)
    bbox = [new_x_min, new_y_min, new_x_max, new_y_max]
    img_crop = crop(img=img, bbox=bbox)

    return img_crop

def crop(img, bbox=None, crop_size=224):
    if bbox is None:
        h, w = img.shape[:-1] #h,w,channel [:-1] beside the final element, such as channel 
        x = random.randint(0, w-crop_size[1])
        y = random.randint(0, h-crop_size[0])
        if isinstance(crop_size, int):
            crop_size = (crop_size, crop_size)  # h, w
        x_min, y_min, x_max, y_max = x, y, x+crop_size[1], y+crop_size[0]
    else:
        x_min, y_min, x_max, y_max = bbox

    crop_img = img[y_min:y_max, x_min:x_max]

    return crop_img

def box_fix_crop_v2(img, bbox, scale=0., c1=0.7, c2=0.6, crop_size=299):

    height, width, _ = img.shape
    if isinstance(crop_size, int):
        crop_size = (crop_size, crop_size)  # h, w

    xmin, ymin, xmax, ymax = bbox
    center_x = (xmax + xmin) / 2.
    center_y = (ymax + ymin) / 2.
    bbox_w, bbox_h = xmax-xmin, ymax-ymin
    
    # 取边框最大值为方框尺寸
    if bbox_h > bbox_w:
        y_side_ = bbox_h
    else:
        y_side_ = bbox_w
    # 对大边框减小边框宽度
    if c1 is not None and c2 is not None:
        if y_side_ >= c1 * width:
            y_side_ = int(c2 * y_side_)
    # 扩大方框尺寸
    y_side_ = int(y_side_ * (1 + scale))
    # 保证方框边长小于图片尺寸
    if y_side_ > height or y_side_ > width:
        y_side_ = min(height, width)
    # 如果方框比较小，设置为稍微大点的值
    if y_side_ < crop_size[0] * (1 + scale):
        y_side_ = crop_size[0] * (1 + scale)

    x_side = y_side_
    y_side = y_side_
    crop_info = [int(center_y - y_side / 2), int(center_y + y_side / 2),
                    int(center_x - x_side / 2), int(center_x + x_side / 2)]
    if crop_info[0] <= 0:
        crop_info[0] = 0
        crop_info[1] = y_side
    if crop_info[1] >= height:
        crop_info[1] = height
        crop_info[0] = max(0, height - y_side)
    if crop_info[2] <= 0:
        crop_info[2] = 0
        crop_info[3] = x_side
    if crop_info[3] >= width:
        crop_info[3] = width
        crop_info[2] = max(0, width - x_side)
    
    bbox = [int(crop_info[2]), int(crop_info[0]), int(crop_info[3]), int(crop_info[1])]
    img_crop = crop(img=img, bbox=bbox)

    return img_crop

class T7_Box_Dataset(Dataset):
    '''
    T2_Box_Dataset
    '''
    def __init__(self, dataframe, code2label, img_size=(224, 224), transform=None):
        self.df = dataframe
        if transform is None:
            self.seg_transforms = resize_transform_basic()
        else:
            self.seg_transforms = transform

        self.code2label = code2label
        self.img_size = img_size

    def __len__(self):
        return self.df.shape[0]

    def __getitem__(self, idx):

        imgPath = self.df["image"].iloc[idx]
        code = str(self.df["code"].iloc[idx])
        try:
            image = cv2.imread(imgPath)
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        except:
            print(imgPath)
            raise ValueError("The image path is incorrect, please check the image path.")

        if str(self.df["box"].iloc[idx]) == '0':
            image = crop(img=image, bbox=None, crop_size=self.img_size)
        else:
            box_tree = read_xml(self.df["box"].iloc[idx])
            
            x_min = int(find_nodes(box_tree, "object/bndbox/xmin")[0].text)
            y_min = int(find_nodes(box_tree, "object/bndbox/ymin")[0].text)
            x_max = int(find_nodes(box_tree, "object/bndbox/xmax")[0].text)
            y_max = int(find_nodes(box_tree, "object/bndbox/ymax")[0].text)

            # image = box_fix_crop(img=image, bbox=[x_min, y_min, x_max, y_max], size=self.img_size)
            image = box_fix_crop_v2(img=image, bbox=[x_min, y_min, x_max, y_max], scale=0., c1=0.7, c2=0.6, crop_size=self.img_size)
            # image = crop(img=image, bbox=[x_min, y_min, x_max, y_max])

        aug = self.seg_transforms(image=image)
        return aug['image'], [self.code2label[code]], imgPath


class T7_Box_Test_Dataset(Dataset):
    '''
    T2_Box_Dataset
    '''
    def __init__(self, dataframe, code2label, img_size=(224, 224), transform=None):
        self.df = dataframe
        self.img_size = img_size
        if transform is None:
            self.seg_transforms = resize_transform_basic()
        else:
            self.seg_transforms = transform

        self.code2label = code2label

    def __len__(self):
        return self.df.shape[0]

    def __getitem__(self, idx):

        imgPath = self.df["image"].iloc[idx]
        code = self.df["code"].iloc[idx]
        box_tree = read_xml(self.df["box"].iloc[idx])
        try:
            image = cv2.imread(imgPath)
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        except:
            raise ValueError("The image path is incorrect, please check the image path.")
        if self.df["box"].iloc[idx] == '0':
            image = crop(img=image, bbox=None, crop_size=self.img_size)
        else:
            box_tree = read_xml(self.df["box"].iloc[idx])
            
            x_min = int(find_nodes(box_tree, "object/bndbox/xmin")[0].text)
            y_min = int(find_nodes(box_tree, "object/bndbox/ymin")[0].text)
            x_max = int(find_nodes(box_tree, "object/bndbox/xmax")[0].text)
            y_max = int(find_nodes(box_tree, "object/bndbox/ymax")[0].text)

            # image = box_fix_crop(img=image, bbox=[x_min, y_min, x_max, y_max], size=self.img_size)
            image = box_fix_crop_v2(img=image, bbox=[x_min, y_min, x_max, y_max], scale=0., c1=0.7, c2=0.6, crop_size=self.img_size)
        aug = self.seg_transforms(image=image)
        return aug['image'], [self.code2label[code]], imgPath, (x_min+x_max)//2, (y_min+y_max)//2, x_max-x_min, y_max-y_min
